{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "448bda69",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sun Aug 27 06:47:49 2023       \r\n",
      "+-----------------------------------------------------------------------------+\r\n",
      "| NVIDIA-SMI 525.85.12    Driver Version: 525.85.12    CUDA Version: 12.0     |\r\n",
      "|-------------------------------+----------------------+----------------------+\r\n",
      "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\r\n",
      "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\r\n",
      "|                               |                      |               MIG M. |\r\n",
      "|===============================+======================+======================|\r\n",
      "|   0  NVIDIA A100 80G...  On   | 00000001:00:00.0 Off |                    0 |\r\n",
      "| N/A   32C    P0    42W / 300W |      4MiB / 81920MiB |      0%      Default |\r\n",
      "|                               |                      |             Disabled |\r\n",
      "+-------------------------------+----------------------+----------------------+\r\n",
      "                                                                               \r\n",
      "+-----------------------------------------------------------------------------+\r\n",
      "| Processes:                                                                  |\r\n",
      "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\r\n",
      "|        ID   ID                                                   Usage      |\r\n",
      "|=============================================================================|\r\n",
      "|  No running processes found                                                 |\r\n",
      "+-----------------------------------------------------------------------------+\r\n"
     ]
    }
   ],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7c15195d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0'\n",
    "os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ec24e3cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, AutoModelForCausalLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3f146124",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained('meta-llama/Llama-2-13b-hf', trust_remote_code=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "196503d3",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "84237012e7be47d395cd8af4285c9bbe",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)lve/main/config.json:   0%|          | 0.00/625 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2023-08-27 06:47:58,792] [INFO] [real_accelerator.py:133:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-08-27 06:47:59.161862: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2023-08-27 06:47:59.813193: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "96913fc1fb684d1fa39c49b598db3c4f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "215e08b12a73476a99f46d3f8859b794",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)of-00003.safetensors:   0%|          | 0.00/9.95G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8589bbd794b04ac295a44bcc1c1ef35d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)of-00003.safetensors:   0%|          | 0.00/9.90G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e30a5313a51c4665af7eb577f3a7a84c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)of-00003.safetensors:   0%|          | 0.00/6.18G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d87d56c5a8de49689abc18673801ae29",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "191ded3c03b948c5aca0bfcdac8e95bc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)neration_config.json:   0%|          | 0.00/183 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "model_meso = AutoModelForCausalLM.from_pretrained(\n",
    "    'mesolitica/llama-13b-hf-2048-fpf', \n",
    "    device_map=\"auto\", \n",
    "    load_in_4bit=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "56f369c0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c7d028fa88414863a11b97cd496d867b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "model_base = AutoModelForCausalLM.from_pretrained(\n",
    "    'meta-llama/Llama-2-13b-hf', \n",
    "    device_map=\"auto\", \n",
    "    load_in_4bit=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "95cac29f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import StoppingCriteria, StoppingCriteriaList"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "361c6967",
   "metadata": {},
   "outputs": [],
   "source": [
    "class StoppingCriteriaSub(StoppingCriteria):\n",
    "\n",
    "    def __init__(self, stops = [], encounters=1):\n",
    "        super().__init__()\n",
    "        self.stops = [stop.to(\"cuda\") for stop in stops]\n",
    "\n",
    "    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor):\n",
    "        for stop in self.stops:\n",
    "            if torch.all((stop == input_ids[0][-len(stop):])).item():\n",
    "                return True\n",
    "\n",
    "        return False\n",
    "    \n",
    "stop_words = [\"\\n\", \"\\nYou:\", \"You:\"]\n",
    "stop_words_ids = [tokenizer(stop_word, return_tensors='pt')['input_ids'].squeeze()[1:] for stop_word in stop_words]\n",
    "stopping_criteria = StoppingCriteriaList([StoppingCriteriaSub(stops=stop_words_ids)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "5f512499",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompts = [\n",
    "    'Continue the chat dialogue below. Write a single reply for the character \"Assistant\". You can ask it anything you want and it will do its best to give you accurate and relevant information.',\n",
    "    'Teruskan dialog di bawah. Tulis satu jawapan untuk watak \"Assistant\". Anda boleh bertanya kepadanya apa sahaja yang anda mahu dan ia akan melakukan yang terbaik untuk memberi anda maklumat yang tepat dan berkaitan.'\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "ac26c716",
   "metadata": {},
   "outputs": [],
   "source": [
    "questions = [\n",
    "    'siapa perdana menteri malaysia',\n",
    "    'KWSP tu apa',\n",
    "    'awat rafizi introduced vending machine?',\n",
    "    'kerajaan madani banyak membantu tak?'\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "92578a1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "max_new_tokens = 500\n",
    "temperature = 0.8\n",
    "top_p = 0.95\n",
    "top_k = 20\n",
    "num_beams = 1\n",
    "do_sample = True\n",
    "repetition_penalty = 1.15"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "2324b615",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Continue the chat dialogue below. Write a single reply for the character \"Assistant\". You can ask it anything you want and it will do its best to give you accurate and relevant information.\n",
      "You: siapa perdana menteri malaysia\n",
      "Assistant:\n",
      "Tun Dr Mahathir bin Mohamad adalah Perdana Menteri Malaysia ke-4 dan berkhidmat dari 1981 hingga 2003. Beliau telah mengetuai Barisan Nasional (BN) untuk memenangi pilihan raya umum dengan majoriti besar sejak tiga kali berturut-turut iaitu pada tahun 1995, 1999 dan 2004. Beliau juga merupakan Perdana Menteri pertama yang dipilih oleh Majlis Tindakan Ekonomi Negara (NEAC), satu agensi kerajaan baru yang ditubuhkan di bawah pentadbirannya. Beliau berjaya mengurus Malaysia sebagai salah sebuah negara paling teratur di Asia. Beliau telah dilantik sebagai Pengerusi Kelab Perkembangan Pemimpin Dunia anjuran World Economic Forum dan dinobatkan sebagai 'The Asian of The Year' oleh majalah Newsweek atas kejayaan beliau dalam menstabilkan ekonomi Malaysia serta membawa pelaburan langsung asing (FDI).</s>\n",
      "\n",
      "Continue the chat dialogue below. Write a single reply for the character \"Assistant\". You can ask it anything you want and it will do its best to give you accurate and relevant information.\n",
      "You: KWSP tu apa\n",
      "Assistant:\n",
      "EPF is an acronym for Employees Provident Fund, which is a retirement fund in Malaysia that provides financial benefits and other assistance to its members who have reached age 55 or older. The fund was established in 1962 as the Central Provident Fund and became known as the Employee Provident Fund (EPF) after being merged with the Retirement Benefits Act 1967. It has grown to become one of the largest pension funds in Asia, providing social security and economic benefits to millions of individuals across various industries. \n",
      "\n",
      "Continue the chat dialogue below. Write a single reply for the character \"Assistant\". You can ask it anything you want and it will do its best to give you accurate and relevant information.\n",
      "You: awat rafizi introduced vending machine?\n",
      "Assistant:\n",
      "he is an entrepreneur, i am surprised that he knows this much about finance. He has been on our platform since the beginning of time! \n",
      "\n",
      "Continue the chat dialogue below. Write a single reply for the character \"Assistant\". You can ask it anything you want and it will do its best to give you accurate and relevant information.\n",
      "You: kerajaan madani banyak membantu tak?\n",
      "Assistant:\n",
      "Sebenarnya kerajaan madani ini sebenarnya adalah kesinambungan daripada kerajaan terdahulu iaitu kerajaan Pakatan Harapan dan Perikatan Nasional yang telah memerintah Malaysia dalam dua penggal yang lalu. Kedua-dua parti politik ini dilihat lebih cenderung untuk memberikan perkhidmatan serta menjalankan tanggungjawab mereka untuk semua rakyat Malaysia, tanpa mengenali latar belakang agama mahupun kaum seseorang. Dari segi ekonomi, kerajaan ini telah berjaya membawa negara ke tahap yang lebih baik dan stabil, di mana Keluaran Dalam Negeri Kasar (KDNK) Malaysia telah mencapai RM1.5 trilion setakat suku pertama 2023. Ini merupakan satu petanda jelas bahawa ekonomi Malaysia sedang menuju ke arah yang betul dengan pembentukan kerajaan baru, dengan gabungan parti-parti politik yang berbeza ideologi yang mempunyai matlamat untuk bekerja demi masa depan dan kebajikan seluruh rakyat Malaysia. Kerajaan Madani telah dilancarkan secara rasmi oleh Yang Amat Berhormat Perdana Menteri Malaysia, Datuk Seri Anwar Ibrahim pada 19 November 2022. Antara objektif utama kerajaan ini adalah untuk mewujudkan masyarakat yang harmoni, inklusif, dan progresif. Selain itu, kerajaan juga ingin membina hubungan erat dengan negara luar, terutamanya dalam bidang ekonomi dan perdagangan. Kerajaan Madani turut menekankan aspek integriti dalam pentadbirannya, di mana ahli-ahli Parlim\n",
      "\n",
      "Teruskan dialog di bawah. Tulis satu jawapan untuk watak \"Assistant\". Anda boleh bertanya kepadanya apa sahaja yang anda mahu dan ia akan melakukan yang terbaik untuk memberi anda maklumat yang tepat dan berkaitan.\n",
      "You: siapa perdana menteri malaysia\n",
      "Assistant:\n",
      "Perdana Menteri Malaysia adalah Tun Dr Mahathir bin Mohamad. \n",
      "\n",
      "Teruskan dialog di bawah. Tulis satu jawapan untuk watak \"Assistant\". Anda boleh bertanya kepadanya apa sahaja yang anda mahu dan ia akan melakukan yang terbaik untuk memberi anda maklumat yang tepat dan berkaitan.\n",
      "You: KWSP tu apa\n",
      "Assistant:\n",
      "Pihak Kerajaan telah menubuhkan sebuah tabung simpanan persaraan, dikenali sebagai 'Kumpulan Wang Simpanan Pekerja', atau lebih dikenali dengan nama singkatnya (KWSP). Ia membolehkan ahli-ahli membuat caruman secara sukarela untuk menyimpan pendapatan pekerjaan mereka demi hari tua mereka. \n",
      "\n",
      "Teruskan dialog di bawah. Tulis satu jawapan untuk watak \"Assistant\". Anda boleh bertanya kepadanya apa sahaja yang anda mahu dan ia akan melakukan yang terbaik untuk memberi anda maklumat yang tepat dan berkaitan.\n",
      "You: awat rafizi introduced vending machine?\n",
      "Assistant:\n",
      "the idea was to reduce cost, increase revenue, and make everyone happy. \n",
      "\n",
      "Teruskan dialog di bawah. Tulis satu jawapan untuk watak \"Assistant\". Anda boleh bertanya kepadanya apa sahaja yang anda mahu dan ia akan melakukan yang terbaik untuk memberi anda maklumat yang tepat dan berkaitan.\n",
      "You: kerajaan madani banyak membantu tak?\n",
      "Assistant:\n",
      "Kerajaan Madani telah menjadi sebuah kerajaan yang dapat memperbaiki keadaan politik, ekonomi, dan juga sosial dalam negara.</s>\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for p in prompts:\n",
    "    for q in questions:\n",
    "        query = f\"\"\"\n",
    "{p}\n",
    "You: {q}\n",
    "Assistant:\n",
    "        \"\"\"\n",
    "        query = query.strip()\n",
    "        inputs = tokenizer(query, return_tensors='pt')\n",
    "        input_ids = inputs['input_ids'].to(model_meso.device)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            generation_output = model_meso.generate(\n",
    "                input_ids=input_ids,\n",
    "                return_dict_in_generate=True,\n",
    "                output_scores=True,\n",
    "                max_new_tokens=max_new_tokens,\n",
    "                temperature=temperature,\n",
    "                top_p=top_p,\n",
    "                top_k=top_k,\n",
    "                num_beams=num_beams,\n",
    "                do_sample=do_sample,\n",
    "                stopping_criteria=stopping_criteria,\n",
    "                repetition_penalty=repetition_penalty,\n",
    "        )\n",
    "        t = tokenizer.decode(generation_output.sequences[0])\n",
    "        t = t.split('Assistant: ')[1]\n",
    "        t = t.split('You:')[0]\n",
    "        print(query)\n",
    "        print(t)\n",
    "        print()\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "e578e820",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Continue the chat dialogue below. Write a single reply for the character \"Assistant\". You can ask it anything you want and it will do its best to give you accurate and relevant information.\n",
      "You: siapa perdana menteri malaysia\n",
      "Assistant:\n",
      "The Prime Minister of Malaysia is Tun Dr Mahathir Mohamad, who took office in May 2018.\n",
      "\n",
      "\n",
      "Continue the chat dialogue below. Write a single reply for the character \"Assistant\". You can ask it anything you want and it will do its best to give you accurate and relevant information.\n",
      "You: KWSP tu apa\n",
      "Assistant:\n",
      "I am an AI-powered virtual assistant created by KWSP, one of Malaysia's leading technology companies specializing in data analysis and processing solutions for businesses. My role is to provide assistance to customers regarding their queries related to our products and services.\n",
      "\n",
      "\n",
      "Continue the chat dialogue below. Write a single reply for the character \"Assistant\". You can ask it anything you want and it will do its best to give you accurate and relevant information.\n",
      "You: awat rafizi introduced vending machine?\n",
      "Assistant:\n",
      "I'm sorry, but I don't understand your question. Please try again.\n",
      "\n",
      "\n",
      "Continue the chat dialogue below. Write a single reply for the character \"Assistant\". You can ask it anything you want and it will do its best to give you accurate and relevant information.\n",
      "You: kerajaan madani banyak membantu tak?\n",
      "Assistant:\n",
      "Kerajaan Madani adalah sebuah kerajaan yang berbasis di negeri Malaysia, dan menyediakan pembantu serta pelayan dari tingkat kebawah hingga tertinggi.\n",
      "\n",
      "\n",
      "Teruskan dialog di bawah. Tulis satu jawapan untuk watak \"Assistant\". Anda boleh bertanya kepadanya apa sahaja yang anda mahu dan ia akan melakukan yang terbaik untuk memberi anda maklumat yang tepat dan berkaitan.\n",
      "You: siapa perdana menteri malaysia\n",
      "Assistant:\n",
      "Yang Berkhidmat Datuk Seri Dr Mahathir bin Mohamad. Dengan kepada tindakan, pembentung perlembagaan Negara Malaysia telah ditetapkan sebagai Yang Dipertuan Agong. Kebanyakan rakyat berpendapat penubuhan perlembagaan baru tidak cukup menyelesaikan isu-isu politik negara. Perdana Menteri dan Timbalan Perdana Menteri terpilih oleh Majlis Legislative (Dewan Rakyat) bersamaan dengan pemerintahan yang dipimpinnya. Menurut Perlembagaan, Perdana Menteri tidak dapat dipecat tanpa alasan sebelum diberhentikan oleh Yang Dipertuan Agong atau membubarkan Parlimen, atau oleh kelompok memuncaki keputusan Yang Dipertuan Agong, atau memecat oleh Yang Dipertuan Agong karena tuduhan korupsi serius.\n",
      "Bagi lebih lanjut maksud semua itu, sila rujuk ke artikel di dalam Wikipedia ini: https://en.wikipedia.org/wiki/Prime_Minister_of_Malaysia\n",
      "Perdana Menteri Malaysia\n",
      "Awal pekerjaan\n",
      "1 Tunku Abdul Rahman (1957–1963)\n",
      "2 Tun Razak (1970–1976)\n",
      "3 Hussein Onn (1976–1981)\n",
      "4 Mahathir (1981–2003)\n",
      "5 Abdullah Ahmad Badawi (2003–2009)\n",
      "6 Najib (2009–sekarang)\n",
      "Majority leader of the House\n",
      "Deputy Prime Minister of Malaysia\n",
      "Cabinet of Malaysia\n",
      "Malaysian federal budget\n",
      "Leader of the Government Business in Parliament\n",
      "Retrieved from \"https://simple.wikipedia.org/w/index.php?title=User:Tonydutch\n",
      "\n",
      "Teruskan dialog di bawah. Tulis satu jawapan untuk watak \"Assistant\". Anda boleh bertanya kepadanya apa sahaja yang anda mahu dan ia akan melakukan yang terbaik untuk memberi anda maklumat yang tepat dan berkaitan.\n",
      "You: KWSP tu apa\n",
      "Assistant:\n",
      "Warga Malaysia yang telah menyambung perkhidmatannya atau sudah membayar sisa sepanjang tiga (3) bulan tidak berlaku hutang dengan Kumpulan Wang Simpanan Pekerja(KWSP). Perkhidmatan itu adalah program pembayaran sisa sepanjang tiga (3) bulan yang mula diterima oleh pelabur KWSP. Ia menerima manfaat daripada pengecualian penyampaian dan kenaikan gaji pendapatan RM500 per bulan untuk semua orang yang menyambung perkhidmatannya, sebagaimana disenaraikan dalam Peruntukan Akaun 2019, sebagai tambahan ke atas RM500 yang diperolehi pada tahun 2018. Manfaat yang ditawarkan terhadap penyampaian dan kenaikan gaji pendapatan, juga membolehkan warga Malaysia yang berpaya tinggi dapat menggunakan lebih banyak untuk membeli barang-barangan.\n",
      "\n",
      "\n",
      "Teruskan dialog di bawah. Tulis satu jawapan untuk watak \"Assistant\". Anda boleh bertanya kepadanya apa sahaja yang anda mahu dan ia akan melakukan yang terbaik untuk memberi anda maklumat yang tepat dan berkaitan.\n",
      "You: awat rafizi introduced vending machine?\n",
      "Assistant:\n",
      "Yes, the vending machines are installed to provide foods for all staffs in our department. They have been working very well so far.\n",
      "\n",
      "\n",
      "Teruskan dialog di bawah. Tulis satu jawapan untuk watak \"Assistant\". Anda boleh bertanya kepadanya apa sahaja yang anda mahu dan ia akan melakukan yang terbaik untuk memberi anda maklumat yang tepat dan berkaitan.\n",
      "You: kerajaan madani banyak membantu tak?\n",
      "Assistant:\n",
      "Yup, I've been working with them for many years now... It is an honor to work with such a reputable company!</s>\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for p in prompts:\n",
    "    for q in questions:\n",
    "        query = f\"\"\"\n",
    "{p}\n",
    "You: {q}\n",
    "Assistant:\n",
    "        \"\"\"\n",
    "        query = query.strip()\n",
    "        inputs = tokenizer(query, return_tensors='pt')\n",
    "        input_ids = inputs['input_ids'].to(model_base.device)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            generation_output = model_base.generate(\n",
    "                input_ids=input_ids,\n",
    "                return_dict_in_generate=True,\n",
    "                output_scores=True,\n",
    "                max_new_tokens=max_new_tokens,\n",
    "                temperature=temperature,\n",
    "                top_p=top_p,\n",
    "                top_k=top_k,\n",
    "                num_beams=num_beams,\n",
    "                do_sample=do_sample,\n",
    "                stopping_criteria=stopping_criteria,\n",
    "                repetition_penalty=repetition_penalty,\n",
    "        )\n",
    "        t = tokenizer.decode(generation_output.sequences[0])\n",
    "        t = t.split('Assistant: ')[1]\n",
    "        t = t.split('You:')[0]\n",
    "        print(query)\n",
    "        print(t)\n",
    "        print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0829415e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
